# -*- coding: utf-8 -*-
# FileName:     detector_gui.py
# time:         23/2/8 008 上午 11:41
# Author:       Zhou Hang
# Description:  基于YOLOv5的前端ui开发
import sys
import time
from pathlib import Path

import cv2
from PyQt5.QtCore import QTimer, QThread, pyqtSignal
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtWidgets import QApplication, QMainWindow, QMessageBox, QFileDialog
from qt_material import apply_stylesheet
from Yolo.model import YoloModel
from ui.detector_ui import Ui_detector_wgt

sys.path.insert(0, './Yolo')

DEFAULT_WEIGHT_PATH = "./Yolo/weight/yolov5s.pt"
CAMERA_ID = 0  # 摄像头编号
SAVE_VIDEO_FPS = 5  # 保存为视频后，视频的FPS
VIDEO_PATH = "./output/videos/"


class InitYolo_t(QThread):
    signal = pyqtSignal(YoloModel)

    def __init__(self, yolo_init_args):
        super(InitYolo_t, self).__init__()
        self.yolo_init_args = yolo_init_args

    def run(self) -> None:
        self.signal.emit(YoloModel(**self.yolo_init_args))


class OpenCamera_t(QThread):
    signal = pyqtSignal(bool)

    def __init__(self):
        super(OpenCamera_t, self).__init__()

    def run(self) -> None:
        pass


class DetectorGui(QMainWindow, Ui_detector_wgt):
    def __init__(self):
        super(DetectorGui, self).__init__()
        self.setupUi(self)
        # 界面状态
        self.side_fold_state = False
        self.parameter_wgt_width = 262
        # 模型状态
        self.yolo_model_t = None
        self.yolo_model = None
        self.weight_path = DEFAULT_WEIGHT_PATH
        # 摄像头相关
        self.camera_is_open = False
        self.cap = None
        self.vid_writer = None
        self.video_timer = QTimer()
        # 初始化
        self.init_common()
        self.init_video()

    # START COMMON
    def init_common(self):
        self.device_comb.addItem('GPU', '')
        self.device_comb.addItem('CPU', 'cpu')
        self.choose_weight_btn.clicked.connect(self.choose_weight)
        self.open_conf_btn.clicked.connect(self.open_conf)
        self.conf_thres_hs.valueChanged.connect(self.conf_thres_ch)
        self.iou_thres_hs.valueChanged.connect(self.iou_thres_ch)
        self.init_model_btn.clicked.connect(self.init_model)

    def conf_thres_ch(self):
        """
        置信度阈值滑动条
        """
        self.conf_thres_lbl.setText(f'置信度 {self.conf_thres_hs.value()}%')

    def iou_thres_ch(self):
        """
        交并比阈值滑动条
        """
        self.iou_thres_lbl.setText(f'IOU    {self.iou_thres_hs.value()}%')

    def open_conf(self):
        """
        导航栏隐藏与打开
        """
        if self.side_fold_state:  # 如果折叠
            self.parameter_wgt.show()
            self.open_conf_btn.setText('配置<<')
            self.parameter_wgt.setFixedWidth(self.parameter_wgt_width)
            self.side_fold_state = False
        else:
            self.parameter_wgt.hide()
            self.open_conf_btn.setText('配置>>')
            # self.parameter_wgt.setFixedWidth(self.btn_wgt_width)
            self.side_fold_state = True

    def choose_weight(self):
        new_weight_path, _ = QFileDialog.getOpenFileName(self,
                                                         '选择权重文件(.pt)',
                                                         './Yolo/weight/')
        if not new_weight_path:
            QMessageBox.warning(self,
                                f"警告", f"打开权重失败, 当前权重{self.weight_path}",
                                buttons=QMessageBox.Ok,
                                defaultButton=QMessageBox.Ok)
        else:
            self.weight_path = new_weight_path
            print('权重文件地址: ' + str(self.weight_path))

    def init_model(self):
        print("开始在另一个线程中初始化模型...")
        yolo_init_args = {
            "weight_path": self.weight_path,
            "conf_thres": self.conf_thres_hs.value() / 100,
            "iou_thres": self.iou_thres_hs.value() / 100,
            "device": self.device_comb.currentData()
        }
        self.yolo_model_t = InitYolo_t(yolo_init_args)
        self.yolo_model_t.signal.connect(self.init_yolo)
        self.yolo_model_t.start()

    def init_yolo(self, data: YoloModel):
        self.yolo_model = data
        # 加载完成的提示框
        print("模型加载完成...")
        QMessageBox.information(self, "提示", "模型加载完成!")
        self.open_camera_btn.setEnabled(True)
        self.FPS_setting_sb.setEnabled(True)
    # END COMMON

    # START CAMERA
    def init_video(self):
        """
        初始化摄像头捕获视频功能
        """
        self.open_camera_btn.clicked.connect(self.open_camera)
        self.video_timer.timeout.connect(self.show_frame)
        self.open_camera_btn.setEnabled(False)
        self.FPS_setting_sb.setEnabled(False)
        self.FPS_setting_sb.valueChanged.connect(self.FPS_change)

    def FPS_change(self):
        """
        修改实时显示FPS的业务逻辑
        """
        if self.camera_is_open:
            self.video_timer.stop()
            self.video_timer.start(self.get_time_slice())

    @staticmethod
    def get_unrepeatable_name():
        """
        通过时间生成一个不可能重复的文件名
        """
        now = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime(time.time()))
        save_path = VIDEO_PATH + now + '.mp4'
        return save_path

    def get_time_slice(self):
        """
        将fps转化为每多少毫秒显示一帧
        """
        return 1000 // (self.FPS_setting_sb.value() + 10)

    def get_camera_info(self):
        """
        获得摄像头信息
        """
        fps = self.cap.get(cv2.CAP_PROP_FPS)
        w = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        h = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        return fps, w, h

    def open_camera(self):
        if not self.camera_is_open:
            print("正在打开摄像头...")
            self.camera_is_open = True
            self.cap = cv2.VideoCapture(CAMERA_ID)
            self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M', 'J', 'P', 'G'))
            # self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 2)

            if not self.cap.isOpened():
                QMessageBox.warning(self, "警告", "摄像头打开失败!")
            else:
                fps, w, h = self.get_camera_info()
                print(f"摄像头参数 fps:{fps} w:{w} h:{h}")
                fps = SAVE_VIDEO_FPS
                if self.video_save_cb.isChecked():
                    save_path = self.get_unrepeatable_name()
                    self.vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, (w, h))
                print(f"当前帧数 {self.FPS_setting_sb.value()}")
                self.video_timer.start(self.get_time_slice())
                self.open_camera_btn.setText("停止检测")
        else:
            self.camera_is_open = False
            self.video_timer.stop()
            self.cap.release()
            self.open_camera_btn.setText("开始检测")

    def show_frame(self):
        """
        视频计时器timeout之后执行
        """
        name_list = []
        read_succ, raw_img = self.cap.read()  # 读取一帧图像
        if read_succ and raw_img is not None:
            raw_img = cv2.flip(raw_img, 180)
            # 检测
            detect_info = self.yolo_model.detect_img(name_list, raw_img)
            # 将获得的实时检测信息渲染到表格上
            self.real_time_table.update_info(detect_info)
            # 如果保存视频, 在这里写入文件
            if self.video_save_cb.isChecked():
                self.vid_writer.write(raw_img)

            # TODO: 这里暂时直接将图片大小写死了,与界面一致
            detect_result = cv2.resize(raw_img, (640, 480))
            # 这里是由于cv2 cap读取到的图片是BGR格式, 转化为RGB格式才能放到QT界面上执行
            detect_result = cv2.cvtColor(detect_result, cv2.COLOR_BGR2RGB)
            showImage = QImage(detect_result.data,
                               detect_result.shape[1],
                               detect_result.shape[0],
                               QImage.Format_RGB888)
            self.main_lbl.setPixmap(QPixmap.fromImage(showImage))
            self.main_lbl.setScaledContents(True)  # 设置图像自适应界面大小
        else:  # 读写结束，释放资源
            print("摄像头读写结束，释放资源")
            self.timer_video.stop()
            self.cap.release()  # 释放video_capture资源
            self.vid_writer.release()  # 释放video_writer资源
            self.main_lbl.clear()
    # END CAMERA


def main():
    app = QApplication(sys.argv)
    detector = DetectorGui()
    extra = {
        'font_size': '15px',
        'font_family': 'Times New Roman'
    }
    # apply_stylesheet(app, theme='dark_pink.xml', extra=extra)
    detector.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()
